#/*************************************************************************
#    > File Name: Class_self.process.py
#    > Author: Yan Wang
#    > Mail: wangyan@imnu.edu.cn
#    > Created Time: Wed 30 Nov 2022 09:22:04 PM CST
# ************************************************************************/
#!/usr/bin/env python
# -*- coding: utf-8 -*-
from Class_Dir import *
from List_Operation import *

class Process(object):
    process={}

    def __init__(self,p):
        self.process["name"]    = p["name"]
        self.process["kind"]    = p["kind"]
        self.process["pt"]      = p["pt"]
        self.process["cms"]     = p["cms"]
        self.process["detector"]= p["detector"]
        self.process["mode"]    = p["mode"]
        self.process["cluster"] = p["cluster"]
        self.process["theoretical_file_name"] = p["theoretical"].format(fname=self.process["kind"])
        self.process["data_style"] = p["data_style"]
        self.process["special_BP"]            = p["special_BP"]

        self.process["running_method"]        = p["running_method"]
        self.process["MG_lhe_event_exist"]    = p["MG_lhe_event_exist"]
        self.process["MG_hepmc_event_exist"]  = p["MG_hepmc_event_exist"]
        self.process["MG_delphes_event_exist"]= p["MG_delphes_event_exist"]

        self.process["event_kind"]            = self.process["kind"]+"_"+self.process["detector"]+"_pt"+self.process["pt"]+"_"+self.process["cms"]+"TeV"
        self.process["process_event_kind"]    = self.process["name"]+"_"+self.process["kind"]+"_"+self.process["detector"]+"_pt"+self.process["pt"]+"_"+self.process["cms"]+"TeV"



    def Basic_Print(self):
        print("	Process Basic Info:")
        print("		name :",self.process["name"])
        print("		kind :",self.process["kind"])
        print("		pt   :",self.process["pt"])
        print("		c.m.s:",self.process["cms"])
        print("		detector:",self.process["detector"])
        print("		mode :",self.process["mode"])
        print("		cluster :",self.process["cluster"])
        print("		running_method : ",self.process["running_method"])
        print("		event_kind : ",self.process["event_kind"])
        print("		process_event_kind : ",self.process["process_event_kind"])




    def Var(self):
        return(self.process)


    def Set_Process(self, param_dir,param_node):
        self.Set_BP_File(param_dir)
        self.Set_Para_Plot()
        self.Set_BP_Folder_Num(param_node["node_num"])
        self.Set_Card(param_dir)
        return



    def Set_BP_File(self, param_dir):
        self.process["theoretical"]=param_dir["input_folder"] / "data" / self.process["theoretical_file_name"]
        self.process["BP_df"]=Load_THDM_Param(self.process["theoretical"])
        #print(self.process["BP_df"].columns)
        self.process["BP_length"]=len(self.process["BP_df"])
        if self.process["special_BP"]:
            self.process["BP_df"]=self.process["BP_df"].loc[np.array(self.process["special_BP"])-1]
            self.process["BP_df"]=self.process["BP_df"].reset_index(drop=True)
        return

    def BP_File_Print(self):
        print("	The Input BP Param File info is :")
        print("		theoretical  : ",self.process["theoretical"])
        print("		BP_df : ",self.process["BP_df"])
        print("		BP_length : ",self.process["BP_length"])


    def Set_Para_Plot(self):
        plot={}
        if(self.process["kind"]=="mh_mch"):
            plot["df_para"]=["mh","mch"]
            plot["X_name"]="M_{h}"
            plot["Y_name"]="M_{H^{#pm}}"
            plot["bins"]=[10,20,120,9,80,170]
            plot["X_axis"]=self.process["BP_df"]["mh"].to_list()
            plot["Y_axis"]=self.process["BP_df"]["mch"].to_list()
        elif(self.process["kind"]=="sba_tb"):
            plot["df_para"]= ["sba","tb"]
            plot["X_name"]= "sin(#beta-#alpha)"
            plot["Y_name"]= "tan(#beta)"
            plot["bins"]= [10,-0.3,-0.0,12,0,60]
            plot["X_axis"]=self.process["BP_df"]["sba"].to_list()
            plot["Y_axis"]=self.process["BP_df"]["tb"].to_list()
        self.process["plot"]=plot
        return

    def Plot_Print(self):
        plot=self.process["plot"]
        print("	The plot Param is : ")
        print("		df_para : ", plot["df_para"])
        print("		x-name  : ", plot["X_name"])
        print("		y-name  : ", plot["Y_name"])
        print("		bins    : ", plot["bins"])
        print("		x-axis  : ", plot["X_axis"])
        print("		y-axis  : ", plot["Y_axis"])

    def Set_BP_Folder_Num(self, node_num):
        ## 设置集群节点参数，用于并行计算
        if self.process["running_method"]== "ssh":
            # 用ssh方式的时候， 是把bps分割后，送到每一个节点，在节点上再用并行处理， 因此任务数等于节点数
            folder_num=node_num
        elif self.process["running_method"]== "local":
            # local 指直接使用本地机器并行处理，因此任务数等于总bps数量
            folder_num=self.process["BP_length"]
        elif self.process["running_method"]== "slurm":
            # slurm 指直接使用slurm方法提交任务，因此任务数等于总bps数量
            folder_num=self.process["BP_length"]
        elif self.process["running_method"]== "pbs":
            # slurm 指直接使用pbs方法提交任务，因此任务数等于总bps数量
            folder_num=self.process["BP_length"]
        else:
            print("please add correct mg_running_method, and rerun the code. The current value is ", self.process["running_method"])
            exit(0)
        self.process["folder_num"]=folder_num
        if self.process["special_BP"]:
            self.process["folder_num"]=len(self.process["special_BP"])

        return

    def Folder_Num_Print(self):
        print("	The Generated Event Folder Num is :")
        print("		folder_num : ", self.process["folder_num"])


    def Set_Card(self, param_dir):
        param_source = param_dir["input_folder"] / 'card/param_card.dat'
        self.process["param_card"]=param_source

        if self.process["pt"]=="10":
            if self.process["cms"]=="13": 
                self.process["run_card"]=param_dir["input_folder"] / "card/run_card_pt10_13TeV.dat"
            elif self.process["cms"]=="14": 
                self.process["run_card"]=param_dir["input_folder"] / "card/run_card_pt10_14TeV.dat"
            if self.process["detector"]=="ATLAS": 
                    self.process["delphes_card"]=param_dir["input_folder"] / "card/delphes_card_ATLAS_pt10.dat"
            elif self.process["detector"]=="CMS": 
                    self.process["delphes_card"]=param_dir["input_folder"] / "card/delphes_card_CMS_pt10.dat"
        elif self.process["pt"]=="20":
            if self.process["cms"]=="13": 
                self.process["run_card"]=param_dir["input_folder"] / "card/run_card_pt20_13TeV.dat"
            elif self.process["cms"]=="14": 
                self.process["run_card"]=param_dir["input_folder"] / "card/run_card_pt20_14TeV.dat"
            if self.process["detector"]=="ATLAS": 
                self.process["delphes_card"]=param_dir["input_folder"] / "card/delphes_card_ATLAS_pt20.dat"
            elif self.process["detector"]=="CMS": 
                self.process["delphes_card"]=param_dir["input_folder"] / "card/delphes_card_CMS_pt20.dat"
        return

    def Card_Print(self):
        print("	The used MG cards are : ")
        print("		run_card : ", self.process["run_card"])
        print("		param_card : ", self.process["param_card"])
        print("		delphes_card : ", self.process["delphes_card"])



    def Print(self):
        print("\n Process Content is : ")
        self.Basic_Print()
        self.Plot_Print()
        self.Folder_Num_Print()
        self.Card_Print()
        print("_"*50)


